import os
import tensorflow as tf
from tensorflow.keras import layers, Model
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.model_selection import train_test_split

print(tf.__version__)
# ==================== 1. 数据加载与预处理 ====================
# ==================== 加载自定义数据集 ====================
def parse_label(line):
    """解析标签文件（路径 数值）"""
    path, num = line.decode('utf-8').split()
    return path, int(num)

# 读取数据集
dataset_path = "multi_digit_dataset"
labels_df = pd.read_csv(os.path.join(dataset_path, "labels.txt"), sep=' ', header=None, names=['path', 'label'])

# 划分训练集和测试集
train_df, test_df = train_test_split(labels_df, test_size=0.2, random_state=42)

# ==================== 构建数据生成器 ====================
def preprocess_image(image_path, label):
    """加载图像并预处理"""
    # 读取PNG图像（灰度图）
    image = tf.io.read_file(image_path)
    image = tf.image.decode_png(image, channels=1)  # 转为单通道灰度图
    image = tf.image.resize(image, [64, 128])  # 统一尺寸（高64，宽128）
    image = tf.cast(image, tf.float32) / 255.0  # 归一化到[0,1]
    return image, label

# 创建训练和测试数据集
BATCH_SIZE = 32
train_ds = tf.data.Dataset.from_tensor_slices((train_df['path'].values, train_df['label'].values))
train_ds = train_ds.map(preprocess_image).shuffle(1000).batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)

test_ds = tf.data.Dataset.from_tensor_slices((test_df['path'].values, test_df['label'].values))
test_ds = test_ds.map(preprocess_image).batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)

# ==================== 2. 构建CNN模型 ====================
def build_crnn_model(input_shape, num_classes, max_time_steps):
    """构建CRNN模型（显式控制时间步长度）"""
    inputs = layers.Input(shape=input_shape)  # 输入形状：(高, 宽, 通道) = (64, 128, 1)
    
    # ------------------- CNN 特征提取（调整池化层以控制时间步） -------------------
    x = layers.Conv2D(32, (3, 3), activation='relu', padding='same')(inputs)  # (64, 128, 32)
    x = layers.MaxPooling2D((2, 2), padding='same')(x)  # (32, 64, 32)（高减半，宽减半）
    
    x = layers.Conv2D(64, (3, 3), activation='relu', padding='same')(x)  # (32, 64, 64)
    x = layers.MaxPooling2D((2, 2), padding='same')(x)  # (16, 32, 64)（高减半，宽减半）
    
    x = layers.Conv2D(128, (3, 3), activation='relu', padding='same')(x)  # (16, 32, 128)
    x = layers.MaxPooling2D((1, 2), padding='same')(x)  # (16, 16, 128)（高不变，宽减半）
    
    x = layers.Conv2D(256, (3, 3), activation='relu', padding='same')(x)  # (16, 16, 256)
    x = layers.BatchNormalization()(x)
    x = layers.Conv2D(256, (3, 3), activation='relu', padding='same')(x)  # (16, 16, 256)
    
    # ------------------- 关键调整：自适应池化控制时间步 -------------------
    # 目标：将特征图的空间维度调整为 (1, max_time_steps)，使时间步长度等于 max_time_steps
    # 输出形状：(batch_size, 1, max_time_steps, 256)
    x = layers.AdaptiveAveragePooling2D(pool_size=(1, max_time_steps))(x)
    
    # 展平为序列（时间步×特征）：(batch_size, max_time_steps, 256)
    x = layers.Reshape((-1, 256))(x)  # -1 自动计算为 max_time_steps
    
    # ------------------- BiLSTM 序列建模 -------------------
    x = layers.Bidirectional(layers.LSTM(128, return_sequences=True))(x)  # (max_time_steps, 256)
    x = layers.Bidirectional(layers.LSTM(64, return_sequences=True))(x)   # (max_time_steps, 128)
    
    # ------------------- 输出层（每个时间步预测一个数字） -------------------
    outputs = layers.TimeDistributed(layers.Dense(num_classes, activation='softmax'))(x)  # (max_time_steps, 10)
    
    model = Model(inputs=inputs, outputs=outputs)
    return model

# 模型参数
input_shape = (64, 128, 1)  # 图像尺寸（高，宽，通道）
num_classes = 10  # 数字0-9
max_time_steps = 128  # 时间步长度（与CNN输出的特征图尺寸相关）
# 初始化模型
model = build_crnn_model(input_shape, num_classes, max_time_steps)
model.summary()

def preprocess_labels_tf(label_str, max_time_steps, char_to_idx):
    """TensorFlow 版本：将字符串标签转换为索引序列（填充或截断到固定长度）"""
    # 将字符串拆分为单个字符（返回字节张量）
    chars = tf.strings.unicode_split(label_str, 'UTF-8')  # 输出：tf.Tensor(b'1', b'2', b'3') 形式
    
    # 将字符转换为索引（需处理填充和截断）
    # 步骤1：将字符映射到索引（使用字典查找）
    # 注意：char_to_idx 需转换为 TensorFlow 可用的格式（如 lookup table）
    keys = tf.constant(list(char_to_idx.keys()), dtype=tf.string)
    values = tf.constant(list(char_to_idx.values()), dtype=tf.int32)
    table = tf.lookup.StaticHashTable(
        tf.lookup.KeyValueTensorInitializer(keys, values, key_dtype=tf.string, value_dtype=tf.int32),
        default_value=-1  # 默认值（填充位置用-1表示）
    )
    
    # 步骤2：查找每个字符的索引（输出：[1, 2, 3] 形式的张量）
    label_indices = table.lookup(chars)  # 输出：tf.Tensor([1, 2, 3], dtype=int32)
    
    # 步骤3：截断或填充到固定长度 max_time_steps
    # 截断过长的序列
    label_indices = label_indices[:max_time_steps]
    # 步骤3：截断或填充到固定长度 max_time_steps（用-1填充）
    pad_length = max_time_steps - tf.shape(label_indices)[0]
    label_indices = tf.pad(label_indices, [[0, pad_length]], constant_values=-1)  # 关键：填充值为-1
    
    return label_indices

def ctc_loss(y_true, y_pred):
    """修正后的 CTC 损失函数"""
    # 调试：打印形状
    print("y_true shape:", y_true.shape)        # 应输出 (batch_size, max_time_steps)
    print("y_pred shape:", y_pred.shape)        # 应输出 (batch_size, max_time_steps, num_classes)
    print("input_length shape:", input_length.shape)  # 应输出 (batch_size,)
    print("label_length shape:", label_length.shape)  # 应输出 (batch_size,)
    # ------------------- 计算 input_length（模型输出的时间步长度） -------------------
    # 假设模型输出的时间步长度固定为 max_time_steps（如128）
    input_length = tf.ones(tf.shape(y_pred)[0], dtype=tf.int32) * tf.shape(y_pred)[1]  # 形状：(batch_size,)
    
    # ------------------- 计算 label_length（真实标签的实际长度） -------------------
    # 统计每个标签序列中非-1的元素个数（即实际长度）
    label_length = tf.math.count_nonzero(y_true != -1, axis=1, dtype=tf.int32)  # 形状：(batch_size,)
    
    # ------------------- 计算 CTC 损失 -------------------
    # 注意：y_pred 需是 logits（未经过 softmax），若模型输出已用 softmax，需先取 log
    # 若模型输出层使用了 softmax，需注释下一行，并确保 y_pred 是 logits
    y_pred_logits = tf.math.log(y_pred + 1e-8)  # 避免 log(0)
    
    loss = tf.keras.backend.ctc_batch_cost(
        y_true=y_true,
        y_pred=y_pred_logits,
        input_length=input_length,
        label_length=label_length
    )
    return tf.reduce_mean(loss)  # 平均损失
# 编译模型（需自定义损失函数，因序列长度可变）
# 注：实际需根据CRNN的CTC损失调整，此处简化为示例
model.compile(optimizer='adam', loss=ctc_loss, metrics=['accuracy'])

char_to_idx = {str(i): i for i in range(num_classes)}  # 字符到索引的映射（0-9）
def preprocess_image_and_label(image_path, label_str):
    """加载图像并预处理标签（TensorFlow 兼容版本）"""
    # ------------------- 图像预处理 -------------------
    # 读取并解码 PNG 图像（灰度图）
    image = tf.io.read_file(image_path)
    image = tf.image.decode_png(image, channels=1)  # 单通道灰度图
    image = tf.image.resize(image, [64, 128])       # 统一尺寸（高64，宽128）
    image = tf.cast(image, tf.float32) / 255.0      # 归一化到[0,1]
    
    # ------------------- 标签预处理（TensorFlow 版本） -------------------
    # 调用 TensorFlow 版本的标签处理函数
    label_indices = preprocess_labels_tf(
        label_str=label_str,
        max_time_steps=max_time_steps,  # 与模型时间步长度一致（如128）
        char_to_idx=char_to_idx         # 字符到索引的映射字典
    )
    
    return image, label_indices

# 假设 train_df 和 test_df 是包含图像路径和标签字符串的 DataFrame
BATCH_SIZE = 32

# 构建训练数据集
train_ds = tf.data.Dataset.from_tensor_slices((
    train_df['path'].values.astype(str),  # 图像路径（字符串张量）
    train_df['label'].values.astype(str)  # 标签字符串（字符串张量）
))
train_ds = train_ds.map(
    lambda img_path, label_str: preprocess_image_and_label(img_path, label_str),
    num_parallel_calls=tf.data.AUTOTUNE  # 并行处理加速
).shuffle(1000).batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)

# 构建测试数据集
test_ds = tf.data.Dataset.from_tensor_slices((
    test_df['path'].values.astype(str),
    test_df['label'].values.astype(str)
))
test_ds = test_ds.map(
    lambda img_path, label_str: preprocess_image_and_label(img_path, label_str),
    num_parallel_calls=tf.data.AUTOTUNE
).batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)


# ==================== 3. 训练模型 ====================
# 训练（示例，实际需调整超参数）
history = model.fit(
    train_ds,
    validation_data=test_ds,
    epochs=10,
    batch_size=BATCH_SIZE
)

# 评估测试集
test_loss, test_acc = model.evaluate(test_ds)
print(f"测试集准确率: {test_acc:.4f}")

# ==================== 4. 预测演示 ====================
def predict_image(image_path):
    """预测单张图像的数字"""
    image = tf.io.read_file(image_path)
    image = tf.image.decode_png(image, channels=1)
    image = tf.image.resize(image, [64, 128])
    image = tf.cast(image, tf.float32) / 255.0
    image = tf.expand_dims(image, axis=0)  # 增加批次维度
    preds = model.predict(image)  # (1, 时间步, num_classes)
    pred_labels = tf.argmax(preds, axis=-1).numpy()[0]  # 取最大概率的类别
    return pred_labels

# 预测示例
image_path = "multi_digit_dataset/42_0.png"
predicted_digits = predict_image(image_path)
print(f"预测的数字: {predicted_digits}")
